﻿#include <gtest/gtest.h>

#include <algorithm>
#include <iostream>
#include <vector>

#include <xy/geom/nurbs.h>
#include <xy/math/matrix.h>

using namespace xy;
using namespace xy::math;
using namespace xy::geom;

class TestGeomNurbs : public ::testing::Test
{
  protected:
    std::mt19937 m_random_engine;

    void SetUp() override
    {
        m_random_engine.seed(std::random_device{}());
    }

    void TearDown() override {
        // nothing
    };

  public:
    std::vector<vec3> GenerateRandomPoints(int num_points)
    {
        std::uniform_real_distribution<double> dist(-1.0, 1.0);
        std::vector<vec3> points(num_points);
        for (int i = 0; i < num_points; ++i)
        {
            points[i] = vec3(dist(m_random_engine), dist(m_random_engine), dist(m_random_engine));
        }
        return points;
    }

    std::vector<vec2> GenerateRandomPoints2D(int num_points)
    {
        std::uniform_real_distribution<double> dist(-1.0, 1.0);
        std::vector<vec2> points(num_points);
        for (int i = 0; i < num_points; ++i)
        {
            points[i] = vec2(dist(m_random_engine), dist(m_random_engine));
        }
        return points;
    }
};

TEST_F(TestGeomNurbs, horner)
{
    {
        // 1 + 2u + 3u^2 + 4u^3
        std::vector<vec3> cpts = {vec3{1.0, 1.0, 1.0}, vec3{2.0, 2.0, 2.0}, vec3{3.0, 3.0, 3.0}, vec3{4.0, 4.0, 4.0}};
        EXPECT_EQ(vec3(10.0, 10.0, 10.0), horner1(cpts.data(), cpts.size(), 1));
    }

    {
        //    1  v  v^2
        // 1  1  2  3
        // u  4  5  6
        std::vector<vec2> cpts = {vec2{1.0, 1.0}, vec2{2.0, 2.0}, vec2{3.0, 3.0},
                                  vec2{4.0, 4.0}, vec2{5.0, 5.0}, vec2{6.0, 6.0}};
        EXPECT_EQ(vec2(4.5, 4.5), horner2(cpts.data(), 2, 3, 0.5, 0.25));
    }

    // 幂基转 Beizer 控制点
    {
        std::vector<vec3> cpts = {vec3{1.0, 1.0, 1.0}, vec3{2.0, 2.0, 2.0}, vec3{3.0, 3.0, 3.0}, vec3{4.0, 4.0, 4.0}};
        std::vector<vec3> cpts2(cpts.size());
        power_to_bezier(cpts.data(), cpts.size(), cpts2.data());

        for (int i = 0; i < 100; i++)
        {
            double u = i / 99.0;
            EXPECT_EQ(horner1(cpts.data(), cpts.size(), u), bezier_point(cpts2.data(), cpts2.size(), u));
        }
    }

    {
        // Beizer 转幂基矩阵
        matxd M3(4, 4);
        bezier_to_power_matrix(3, M3.data());
        matxd T3(4, 4);
        T3(0, 0) = T3(3, 3) = 1;
        T3(1, 1) = T3(2, 2) = 3;
        T3(1, 0) = T3(3, 2) = -3;
        T3(2, 0) = T3(3, 1) = 3;
        T3(3, 0) = -1;
        T3(2, 1) = -6;
        EXPECT_LE((M3 - T3).norm(), 1e-10);

        matxd M2(3, 3);
        bezier_to_power_matrix(2, M2.data());
        matxd T2(3, 3);
        T2(0, 0) = T2(2, 2) = 1;
        T2(1, 1) = 2;
        T2(1, 0) = T2(2, 1) = -2;
        T2(2, 0) = 1;
        EXPECT_LE((M2 - T2).norm(), 1e-10);

        matxd M1(2, 2);
        bezier_to_power_matrix(1, M1.data());
        matxd T1(2, 2);
        T1(0, 0) = T1(1, 1) = 1;
        T1(1, 0) = -1;
        EXPECT_LE((M1 - T1).norm(), 1e-10);

        // 幂基转 Beizer 基矩阵
        matxd M3I(4, 4);
        power_to_bezier_matrix(3, M3.data(), M3I.data());
        matxd T3I(4, 4);
        T3I(0, 0) = T3I(3, 3) = 1;
        T3I(1, 1) = T3I(2, 2) = 1.0 / 3;
        T3I(1, 0) = T3I(3, 2) = 1;
        T3I(2, 0) = T3I(3, 1) = 1;
        T3I(3, 0) = 1;
        T3I(2, 1) = 2.0 / 3;
        EXPECT_LE((M3I - T3I).norm(), 1e-10);

        matxd M2I(3, 3);
        power_to_bezier_matrix(2, M2.data(), M2I.data());
        matxd T2I(3, 3);
        T2I(0, 0) = T2I(2, 2) = 1;
        T2I(1, 1) = 1.0 / 2;
        T2I(1, 0) = T2I(2, 1) = 1;
        T2I(2, 0) = 1;
        EXPECT_LE((M2I - T2I).norm(), 1e-10);

        matxd M1I(2, 2);
        power_to_bezier_matrix(1, M1.data(), M1I.data());
        matxd T1I(2, 2);
        T1I(0, 0) = T1I(1, 1) = 1;
        T1I(1, 0) = 1;
        EXPECT_LE((M1I - T1I).norm(), 1e-10);
    }
}

TEST_F(TestGeomNurbs, bezier)
{
    EXPECT_EQ(120, factorial(5));
    EXPECT_FLOAT_EQ(10, binomial(5, 3));
    EXPECT_FLOAT_EQ(20, permutation(5, 3));
    EXPECT_FLOAT_EQ(0.263671875, bernstein(5, 3, 0.75));

    int N = 10;
    auto points = GenerateRandomPoints(N);

    int M = 1000;
    for (int i = 0; i < M; ++i)
    {
        double t = static_cast<double>(i) / M;
        EXPECT_EQ(de_casteljau1(points.data(), N, t), bezier_point(points.data(), N, t));
    }

    auto points2 = bezier_elevate(points.data(), N, 2);
    for (int i = 0; i < M; ++i)
    {
        double t = static_cast<double>(i) / M;
        EXPECT_EQ(bezier_point(points2.data(), N + 2, t), bezier_point(points.data(), N, t));
    }
}

TEST_F(TestGeomNurbs, find_span)
{
    {
        std::vector<double> knots({0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5});
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 0.0), 2);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 0.5), 2);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 1.5), 3);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 2.5), 4);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 3.5), 5);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 4.5), 7);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 5.0), 7);
    }

    {
        std::vector<double> knots({0, 0, 0, 1, 2, 2, 4, 4, 5, 5, 5});
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 0.0), 2);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 0.5), 2);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 1.5), 3);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 2.5), 5);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 3.5), 5);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 4.5), 7);
        EXPECT_EQ(find_span(knots.data(), knots.size(), 2, 5.0), 7);
    }
}

TEST_F(TestGeomNurbs, basis_functions)
{
    {
        std::vector<double> knots({0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5});
        int span = find_span(knots.data(), knots.size(), 2, 2.5);
        EXPECT_EQ(span, 4);

        auto basis = basis_funcs(knots.data(), 2, span, 2.5);
        EXPECT_EQ(basis, std::vector<double>({1.0 / 8, 6.0 / 8, 1.0 / 8}));

        std::vector<std::vector<double>> N(3);
        N[0] = std::vector<double>({1});
        N[1] = std::vector<double>({0.5, 0.5});
        N[2] = std::vector<double>({1.0 / 8, 6.0 / 8, 1.0 / 8});
        EXPECT_EQ(all_basis_funcs(knots.data(), 2, span, 2.5), N);

        EXPECT_EQ(one_basis_func(knots.data(), knots.size(), 2, span, 2.5), 1.0 / 8);
        EXPECT_EQ(one_basis_func(knots.data(), knots.size(), 2, span - 1, 2.5), 6.0 / 8);
        EXPECT_EQ(one_basis_func(knots.data(), knots.size(), 2, span - 2, 2.5), 1.0 / 8);
    }

    {
        std::vector<double> knots({0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5});
        int span = find_span(knots.data(), knots.size(), 2, 2.5);
        EXPECT_EQ(span, 4);
        EXPECT_EQ(deriv_basis_funcs(knots.data(), 2, span, 2, 2.5),
                  std::vector<std::vector<double>>({{0.125, 0.75, 0.125}, {-0.5, 0.0, 0.5}, {1.0, -2.0, 1.0}}));
        EXPECT_EQ(deriv_one_basis_func(knots.data(), knots.size(), 2, span, 2, 2.5),
                  std::vector<double>({0.125, 0.5, 1.0}));
    }
}

TEST_F(TestGeomNurbs, curve)
{
    {
        std::vector<vec3> points(8);
        for (int i = 0; i < 8; i++)
            points[i] = vec3{i * 0.5, i * i * 0.8, sin(i)};
        std::vector<double> knots({0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5});

        EXPECT_EQ(curve_point(points.data(), knots.data(), knots.size(), 2, 2.5),
                  points[2] / 8 + points[3] * 6 / 8 + points[4] / 8);
        EXPECT_EQ(curve_point_by_corner_cut(points.data(), knots.data(), knots.size(), 2, 2.5),
                  points[2] / 8 + points[3] * 6 / 8 + points[4] / 8);
    }

    {
        int m = 7;
        int p = 2;
        int n = m - p - 1;
        std::vector<vec3> points(5);
        for (int i = 0; i < 5; i++)
            points[i] = vec3{i * 0.5, i * i * 0.8, sin(i)};
        std::vector<double> knots({0, 0, 0, 0.4, 0.6, 1, 1, 1});

        // 求导后只有 4 个有效的控制点，剩下的为零（空）
        std::vector<vec3> derivPoints(5);
        derivPoints[0] = (points[1] - points[0]) * 5;
        derivPoints[1] = (points[2] - points[1]) * 10 / 3;
        derivPoints[2] = (points[3] - points[2]) * 10 / 3;
        derivPoints[3] = (points[4] - points[3]) * 5;
        auto PK = curve_deriv_cpts(points.data(), knots.data(), knots.size(), 2, 1, 0, n);
        EXPECT_EQ(PK[0], points);
        EXPECT_EQ(PK[1], derivPoints);
    }

    {
        int m = 10;
        int p = 3;
        int n = m - p - 1;
        std::vector<vec3> points(7);
        for (int i = 0; i < 7; i++)
            points[i] = vec3{i * 0.5, i * i * 0.8, sin(i)};
        std::vector<double> knots({0, 0, 0, 0, 1.0 / 3, 2.0 / 3, 2.0 / 3, 1, 1, 1, 1});

        // 求导后只有 6 个有效的控制点，剩下的为零（空）
        std::vector<vec3> derivPoints(7);
        derivPoints[0] = (points[1] - points[0]) * 9;
        derivPoints[1] = (points[2] - points[1]) * 4.5;
        derivPoints[2] = (points[3] - points[2]) * 4.5;
        derivPoints[3] = (points[4] - points[3]) * 4.5;
        derivPoints[4] = (points[5] - points[4]) * 9;
        derivPoints[5] = (points[6] - points[5]) * 9;
        auto PK = curve_deriv_cpts(points.data(), knots.data(), knots.size(), 3, 1, 0, n);
        EXPECT_EQ(PK[0], points);
        EXPECT_EQ(PK[1], derivPoints);
    }

    {
        int m = 10;
        int p = 2;
        int n = m - p - 1;
        std::vector<vec3> points(8);
        for (int i = 0; i < 8; i++)
            points[i] = vec3{i * 0.5, i * i * 0.8, sin(i)};
        std::vector<double> knots({0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5});

        auto CK1 = curve_deriv_alg1(points.data(), knots.data(), knots.size(), 2, 1, 2.5);
        auto CK2 = curve_deriv_alg2(points.data(), knots.data(), knots.size(), 2, 1, 2.5);
        EXPECT_EQ(CK1[1], (points[2] * -0.5 + points[4] * 0.5));
        EXPECT_EQ(CK2[1], (points[2] * -0.5 + points[4] * 0.5));
    }
}

TEST_F(TestGeomNurbs, decompose)
{
    int m = 7;
    int p = 2;
    int n = m - p - 1;
    std::vector<vec3> points(5);
    for (int i = 0; i < 5; i++)
        points[i] = vec3{i * 0.5, i * i * 0.8, sin(i)};
    std::vector<double> knots({0, 0, 0, 0.4, 0.6, 1, 1, 1});

    auto Q = curve_decompose(points.data(), knots.data(), knots.size(), 2);

    // 注意细分后，Bezier 参数区间为 [0, 1]，通过区间等分来映射到原始曲线上
    int N = 10;
    for (int i = 0; i < N; i++)
    {
        double u = 0.0 + (0.4 - 0.0) * i / N;
        double t = 1.0 * i / N;
        EXPECT_EQ(bezier_point(Q[0].data(), Q[0].size(), t),
                  curve_point(points.data(), knots.data(), knots.size(), 2, u));
    }

    for (int i = 0; i < N; i++)
    {
        double u = 0.4 + (0.6 - 0.4) * i / N;
        double t = 1.0 * i / N;
        EXPECT_EQ(bezier_point(Q[1].data(), Q[1].size(), t),
                  curve_point(points.data(), knots.data(), knots.size(), 2, u));
    }

    for (int i = 0; i < N; i++)
    {
        double u = 0.6 + (1.0 - 0.6) * i / N;
        double t = 1.0 * i / N;
        EXPECT_EQ(bezier_point(Q[2].data(), Q[2].size(), t),
                  curve_point(points.data(), knots.data(), knots.size(), 2, u));
    }
}

TEST_F(TestGeomNurbs, surface)
{
    {
        std::vector<vec3> P(5 * 6);
        for (int i = 0; i < 5; i++)
            for (int j = 0; j < 6; j++)
                P[i * 6 + j] = vec3((double)i, sin(i + j), (double)j);

        std::vector<double> uknots({0, 0.1, 0.2, 0.4, 0.6, 1, 2, 3});
        std::vector<double> vknots({0, 0.1, 0.15, 0.2, 0.5, 0.8, 1, 2, 3, 3});

        // 先算两个方向的非零基函数
        int ui = find_span(uknots.data(), uknots.size(), 2, 0.3);
        int vi = find_span(vknots.data(), vknots.size(), 3, 0.6);
        std::vector<double> uB = basis_funcs(uknots.data(), 2, ui, 0.3);
        std::vector<double> vB = basis_funcs(vknots.data(), 3, vi, 0.6);

        // 组合利用定义计算点
        vec3 S;
        for (int i = 0; i < 3; i++)
        {
            vec3 tmp;
            for (int j = 0; j < 4; j++)
                tmp = tmp + P[(ui - 2 + i) * 6 + vi - 3 + j] * vB[j];
            S = S + tmp * uB[i];
        }

        // 计算曲面上的点，与定义计算的结果比较
        EXPECT_EQ(surface_point(P.data(), uknots.data(), vknots.data(), uknots.size(), vknots.size(), 2, 3, 0.3, 0.6),
                  S);
    }

    {
        std::vector<vec3> P2(5 * 6);
        for (int i = 0; i < 5; i++)
            for (int j = 0; j < 6; j++)
                P2[i * 6 + j] = vec3((double)i, sin(i + j), (double)j);

        std::vector<double> uknots({0, 0.1, 0.2, 0.4, 0.4, 1, 2, 3});
        std::vector<double> vknots({0, 0.1, 0.1, 0.2, 0.5, 0.8, 1, 2, 3, 3});

        double bu = 0.2, bv = 1;
        double eu = 0.4, ev = 1;
        for (int i = 0; i < 100; i++)
        {
            for (int j = 0; j < 100; j++)
            {
                double u = bu + (eu - bu) / 99 * i;
                double v = bv + (ev - bv) / 99 * j;

                int ui = find_span(uknots.data(), uknots.size(), 2, u);
                int vi = find_span(vknots.data(), vknots.size(), 3, v);

                auto Dp1 = deriv_basis_funcs(uknots.data(), 2, ui, 2, u);
                auto Dp2 = deriv_basis_funcs(vknots.data(), 3, vi, 2, v);

                // 输入求导后的基函数，计算导数曲面上的点
                auto f = [=](std::vector<double> p1, std::vector<double> p2) -> vec3 {
                    vec3 p;
                    for (int i = 0; i <= 2; i++)
                    {
                        vec3 tmp;
                        for (int j = 0; j <= 3; j++)
                            tmp = tmp + P2[(ui - 2 + i) * 6 + vi - 3 + j] * p2[j];
                        p = p + tmp * p1[i];
                    }
                    return p;
                };

                // 求曲面上的点
                vec3 p =
                    surface_point(P2.data(), uknots.data(), vknots.data(), uknots.size(), vknots.size(), 2, 3, u, v);

                // 求 2 阶导对应的点
                auto P = surface_deriv_alg1(P2.data(), uknots.data(), vknots.data(), uknots.size(), vknots.size(), 2, 3,
                                            2, u, v);

                // 验证所有偏导
                EXPECT_EQ(P[0][0], p);
                EXPECT_EQ(P[0][0], f(Dp1[0], Dp2[0]));
                EXPECT_EQ(P[1][0], f(Dp1[1], Dp2[0]));
                EXPECT_EQ(P[0][1], f(Dp1[0], Dp2[1]));
                EXPECT_EQ(P[1][1], f(Dp1[1], Dp2[1]));
                EXPECT_EQ(P[2][0], f(Dp1[2], Dp2[0]));
                EXPECT_EQ(P[0][2], f(Dp1[0], Dp2[2]));
            }
        }
    }
}